"""
训练任务相关的Pydantic模式
"""
from datetime import datetime
from typing import Dict, List, Optional, Any
from pydantic import BaseModel, Field, ConfigDict

from app.models.training_job import TrainingStatus, TrainingType


class TrainingJobBase(BaseModel):
    """训练任务基础模式"""
    name: str = Field(..., min_length=1, max_length=100, description="训练任务名称")
    description: Optional[str] = Field(None, max_length=1000, description="训练任务描述")
    training_type: TrainingType = Field(..., description="训练类型")
    model_architecture: str = Field(..., max_length=50, description="模型架构")
    total_epochs: int = Field(..., gt=0, description="总训练轮数")
    hyperparameters: Dict[str, Any] = Field(default_factory=dict, description="超参数")
    training_config: Dict[str, Any] = Field(default_factory=dict, description="训练配置")
    resource_requirements: Dict[str, Any] = Field(default_factory=dict, description="资源需求")
    dataset_config: Dict[str, Any] = Field(default_factory=dict, description="数据集配置")
    data_split_config: Dict[str, Any] = Field(default_factory=dict, description="数据分割配置")
    augmentation_config: Optional[Dict[str, Any]] = Field(None, description="数据增强配置")


class TrainingJobCreate(TrainingJobBase):
    """训练任务创建模式"""
    project_id: int = Field(..., description="项目ID")
    parent_job_id: Optional[int] = Field(None, description="父任务ID")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "智慧照明预测模型训练",
                "description": "基于传感器数据的照明亮度预测模型",
                "training_type": "supervised",
                "model_architecture": "transformer",
                "total_epochs": 100,
                "hyperparameters": {
                    "learning_rate": 0.001,
                    "batch_size": 32,
                    "hidden_size": 512,
                    "num_layers": 6
                },
                "training_config": {
                    "optimizer": "adam",
                    "loss_function": "mse",
                    "early_stopping": True,
                    "patience": 10
                },
                "resource_requirements": {
                    "gpu_count": 1,
                    "memory_gb": 16,
                    "cpu_cores": 4
                },
                "dataset_config": {
                    "data_source_ids": [1, 2],
                    "features": ["temperature", "humidity", "light_level"],
                    "target": "brightness_level"
                },
                "data_split_config": {
                    "train_ratio": 0.7,
                    "val_ratio": 0.2,
                    "test_ratio": 0.1
                },
                "project_id": 1
            }
        }
    )


class TrainingJobUpdate(BaseModel):
    """训练任务更新模式"""
    name: Optional[str] = Field(None, min_length=1, max_length=100, description="训练任务名称")
    description: Optional[str] = Field(None, max_length=1000, description="训练任务描述")
    status: Optional[TrainingStatus] = Field(None, description="训练状态")
    hyperparameters: Optional[Dict[str, Any]] = Field(None, description="超参数")
    training_config: Optional[Dict[str, Any]] = Field(None, description="训练配置")


class TrainingJobInDB(TrainingJobBase):
    """数据库中的训练任务模式"""
    id: int
    status: TrainingStatus
    project_id: int
    created_by: int
    parent_job_id: Optional[int]
    progress_percentage: float
    current_epoch: int
    metrics: Dict[str, Any]
    best_metrics: Dict[str, Any]
    worker_node: Optional[str]
    gpu_ids: List[str]
    log_path: Optional[str]
    checkpoint_path: Optional[str]
    estimated_duration: Optional[int]
    actual_duration: Optional[int]
    started_at: Optional[datetime]
    completed_at: Optional[datetime]
    error_message: Optional[str]
    error_traceback: Optional[str]
    created_at: datetime
    updated_at: datetime
    
    model_config = ConfigDict(from_attributes=True)


class TrainingJob(TrainingJobBase):
    """训练任务响应模式"""
    id: int
    status: TrainingStatus
    project_id: int
    created_by: int
    parent_job_id: Optional[int]
    progress_percentage: float
    current_epoch: int
    metrics: Dict[str, Any]
    best_metrics: Dict[str, Any]
    worker_node: Optional[str]
    gpu_ids: List[str]
    log_path: Optional[str]
    checkpoint_path: Optional[str]
    estimated_duration: Optional[int]
    actual_duration: Optional[int]
    started_at: Optional[datetime]
    completed_at: Optional[datetime]
    error_message: Optional[str]
    created_at: datetime
    updated_at: datetime
    
    model_config = ConfigDict(from_attributes=True)


class TrainingJobSummary(BaseModel):
    """训练任务摘要模式"""
    id: int
    name: str
    training_type: TrainingType
    status: TrainingStatus
    progress_percentage: float
    current_epoch: int
    total_epochs: int
    started_at: Optional[datetime]
    estimated_duration: Optional[int]
    created_at: datetime
    
    model_config = ConfigDict(from_attributes=True)


class TrainingJobList(BaseModel):
    """训练任务列表模式"""
    training_jobs: List[TrainingJobSummary]
    total: int
    page: int
    page_size: int


class TrainingJobControl(BaseModel):
    """训练任务控制模式"""
    action: str = Field(..., regex="^(start|pause|resume|stop|cancel)$", description="操作类型")
    reason: Optional[str] = Field(None, description="操作原因")


class TrainingJobMetrics(BaseModel):
    """训练任务指标模式"""
    epoch: int
    loss: float
    accuracy: Optional[float] = None
    val_loss: Optional[float] = None
    val_accuracy: Optional[float] = None
    learning_rate: float
    timestamp: datetime
    additional_metrics: Dict[str, float] = Field(default_factory=dict)


class TrainingJobLogs(BaseModel):
    """训练任务日志模式"""
    logs: List[str] = Field(..., description="日志行")
    total_lines: int = Field(..., description="总日志行数")
    start_line: int = Field(..., description="起始行号")
    end_line: int = Field(..., description="结束行号")


class TrainingJobStats(BaseModel):
    """训练任务统计模式"""
    total_jobs: int
    running_jobs: int
    completed_jobs: int
    failed_jobs: int
    queued_jobs: int
    avg_training_time: float
    success_rate: float
    jobs_by_type: Dict[str, int]
    resource_utilization: Dict[str, float]